import numpy as np
import pytest
from sklearn import config_context

from skfolio.moments import ImpliedCovariance
from skfolio.prior import BlackLitterman, EmpiricalPrior, FactorModel
from skfolio.utils.equations import equations_to_matrix


def test_views_to_matrix(X):
    views = ["AAPL - BBY == 0.03 ", "CVX - KO== 0.04", "MSFT == 0.06 "]
    groups = np.array([X.columns])

    picking_matrix, views, _, _ = equations_to_matrix(
        groups=groups, equations=views, sum_to_one=True
    )
    np.testing.assert_almost_equal(
        picking_matrix,
        np.array(
            [
                [
                    1.0,
                    0.0,
                    0.0,
                    -1.0,
                    0.0,
                    0.0,
                    0.0,
                    0.0,
                    0.0,
                    0.0,
                    0.0,
                    0.0,
                    0.0,
                    0.0,
                    0.0,
                    0.0,
                    0.0,
                    0.0,
                    0.0,
                    0.0,
                ],
                [
                    0.0,
                    0.0,
                    0.0,
                    0.0,
                    1.0,
                    0.0,
                    0.0,
                    0.0,
                    0.0,
                    -1.0,
                    0.0,
                    0.0,
                    0.0,
                    0.0,
                    0.0,
                    0.0,
                    0.0,
                    0.0,
                    0.0,
                    0.0,
                ],
                [
                    0.0,
                    0.0,
                    0.0,
                    0.0,
                    0.0,
                    0.0,
                    0.0,
                    0.0,
                    0.0,
                    0.0,
                    0.0,
                    0.0,
                    1.0,
                    0.0,
                    0.0,
                    0.0,
                    0.0,
                    0.0,
                    0.0,
                    0.0,
                ],
            ]
        ),
    )
    np.testing.assert_almost_equal(views, np.array([0.03, 0.04, 0.06]))


def test_black_litterman(X):
    views = ["AAPL - BBY == 0.03 ", "CVX - KO== 0.04", "MSFT == 0.06 "]
    n_observations = X.shape[0]
    model = BlackLitterman(views=np.array(views), tau=1 / n_observations)
    model.fit(X)
    assert np.all(model.groups_ == np.asarray(X.columns))
    assert np.all(model.views_ == np.array([0.03, 0.04, 0.06]))
    assert model.picking_matrix_.shape == (3, 20)
    res = model.return_distribution_
    assert hash(res)
    assert res.mu.shape == (20,)
    assert res.covariance.shape == (20, 20)
    np.testing.assert_almost_equal(res.returns, np.asarray(X))
    assert res.cholesky is None
    np.testing.assert_almost_equal(
        res.mu,
        np.array(
            [
                0.02644867,
                0.03067605,
                0.02164973,
                0.00928809,
                0.02765341,
                0.01837103,
                0.01600049,
                0.00939478,
                0.01957696,
                0.00585437,
                0.01234366,
                0.01059111,
                0.03200649,
                0.00907145,
                0.01037331,
                0.00785186,
                0.0263905,
                0.01666245,
                0.00770657,
                0.02098158,
            ]
        ),
    )

    np.testing.assert_almost_equal(
        res.covariance,
        np.array(
            [
                [
                    3.36958455e-04,
                    2.83215895e-04,
                    1.56228019e-04,
                    1.75889683e-04,
                    1.22085731e-04,
                    1.38700761e-04,
                    1.47126181e-04,
                    8.11023032e-05,
                    1.42513814e-04,
                    8.01829122e-05,
                    9.68619215e-05,
                    7.98648783e-05,
                    2.15052791e-04,
                    9.97625137e-05,
                    8.61115392e-05,
                    8.56846596e-05,
                    1.37608444e-04,
                    1.33035904e-04,
                    8.21656196e-05,
                    1.06281426e-04,
                ],
                [
                    2.83215895e-04,
                    1.39288506e-03,
                    2.20637645e-04,
                    2.51583699e-04,
                    1.76280474e-04,
                    1.87547788e-04,
                    2.00212149e-04,
                    8.31723300e-05,
                    1.81476730e-04,
                    8.63635656e-05,
                    1.16657507e-04,
                    9.03627314e-05,
                    2.57932095e-04,
                    1.02358608e-04,
                    1.10805178e-04,
                    8.74489219e-05,
                    2.68568733e-04,
                    1.62238509e-04,
                    8.62748168e-05,
                    1.38462994e-04,
                ],
                [
                    1.56228019e-04,
                    2.20637645e-04,
                    3.93041965e-04,
                    1.94999956e-04,
                    2.20767528e-04,
                    2.42402367e-04,
                    1.49288034e-04,
                    8.83532871e-05,
                    3.12688640e-04,
                    9.85503243e-05,
                    9.19820846e-05,
                    9.86451929e-05,
                    1.58958119e-04,
                    9.03141055e-05,
                    1.04095672e-04,
                    8.01562125e-05,
                    2.67054571e-04,
                    1.52271435e-04,
                    6.72958419e-05,
                    1.98387376e-04,
                ],
                [
                    1.75889683e-04,
                    2.51583699e-04,
                    1.94999956e-04,
                    6.15403112e-04,
                    1.47162606e-04,
                    1.68273040e-04,
                    1.96104575e-04,
                    7.41922212e-05,
                    1.72778409e-04,
                    7.94306295e-05,
                    7.66353577e-05,
                    6.82995013e-05,
                    1.53111000e-04,
                    8.80192972e-05,
                    7.46460062e-05,
                    8.14494768e-05,
                    2.06813151e-04,
                    1.25016166e-04,
                    9.30891364e-05,
                    1.27462037e-04,
                ],
                [
                    1.22085731e-04,
                    1.76280474e-04,
                    2.20767528e-04,
                    1.47162606e-04,
                    3.60817266e-04,
                    2.10617996e-04,
                    1.24203206e-04,
                    8.34453166e-05,
                    2.02142277e-04,
                    9.39485178e-05,
                    8.78493224e-05,
                    9.48295289e-05,
                    1.28918711e-04,
                    8.25067070e-05,
                    8.71845732e-05,
                    6.50007767e-05,
                    3.23877738e-04,
                    1.36417816e-04,
                    5.45391833e-05,
                    2.79531990e-04,
                ],
                [
                    1.38700761e-04,
                    1.87547788e-04,
                    2.42402367e-04,
                    1.68273040e-04,
                    2.10617996e-04,
                    4.82869733e-04,
                    1.24699894e-04,
                    7.82059674e-05,
                    2.18173034e-04,
                    9.79324240e-05,
                    8.48704818e-05,
                    8.37893828e-05,
                    1.28399776e-04,
                    7.87550632e-05,
                    9.10274252e-05,
                    7.15203459e-05,
                    2.41363768e-04,
                    1.24911263e-04,
                    5.27462412e-05,
                    1.95714717e-04,
                ],
                [
                    1.47126181e-04,
                    2.00212149e-04,
                    1.49288034e-04,
                    1.96104575e-04,
                    1.24203206e-04,
                    1.24699894e-04,
                    2.35263666e-04,
                    7.52692463e-05,
                    1.41816466e-04,
                    8.26880747e-05,
                    8.81690672e-05,
                    7.61632563e-05,
                    1.52606397e-04,
                    1.00653474e-04,
                    8.12084880e-05,
                    8.22072630e-05,
                    1.01272428e-04,
                    1.24468155e-04,
                    8.57347703e-05,
                    9.68811805e-05,
                ],
                [
                    8.11023032e-05,
                    8.31723300e-05,
                    8.83532871e-05,
                    7.41922212e-05,
                    8.34453166e-05,
                    7.82059674e-05,
                    7.52692463e-05,
                    1.30985801e-04,
                    8.54684096e-05,
                    6.87330872e-05,
                    9.85385376e-05,
                    8.66156451e-05,
                    8.74914106e-05,
                    7.90839398e-05,
                    8.87031919e-05,
                    7.56739267e-05,
                    6.55688622e-05,
                    9.38343175e-05,
                    6.01344712e-05,
                    7.14216054e-05,
                ],
                [
                    1.42513814e-04,
                    1.81476730e-04,
                    3.12688640e-04,
                    1.72778409e-04,
                    2.02142277e-04,
                    2.18173034e-04,
                    1.41816466e-04,
                    8.54684096e-05,
                    3.01288346e-04,
                    9.66296262e-05,
                    8.64601714e-05,
                    9.41168573e-05,
                    1.45795307e-04,
                    8.85816243e-05,
                    1.00943437e-04,
                    7.67240024e-05,
                    2.16800171e-04,
                    1.38853956e-04,
                    6.40983892e-05,
                    1.77534524e-04,
                ],
                [
                    8.01829122e-05,
                    8.63635656e-05,
                    9.85503243e-05,
                    7.94306295e-05,
                    9.39485178e-05,
                    9.79324240e-05,
                    8.26880747e-05,
                    6.87330872e-05,
                    9.66296262e-05,
                    1.34182373e-04,
                    6.66436668e-05,
                    6.54699514e-05,
                    8.58786148e-05,
                    1.00394611e-04,
                    6.65754039e-05,
                    8.42955525e-05,
                    5.75321383e-05,
                    8.11275129e-05,
                    5.74148452e-05,
                    8.46645084e-05,
                ],
                [
                    9.68619215e-05,
                    1.16657507e-04,
                    9.19820846e-05,
                    7.66353577e-05,
                    8.78493224e-05,
                    8.48704818e-05,
                    8.81690672e-05,
                    9.85385376e-05,
                    8.64601714e-05,
                    6.66436668e-05,
                    2.79272627e-04,
                    1.14620817e-04,
                    1.11349186e-04,
                    8.26738986e-05,
                    1.18642453e-04,
                    7.54084960e-05,
                    7.99199098e-05,
                    1.17689173e-04,
                    6.39652089e-05,
                    8.04271042e-05,
                ],
                [
                    7.98648783e-05,
                    9.03627314e-05,
                    9.86451929e-05,
                    6.82995013e-05,
                    9.48295289e-05,
                    8.37893828e-05,
                    7.61632563e-05,
                    8.66156451e-05,
                    9.41168573e-05,
                    6.54699514e-05,
                    1.14620817e-04,
                    1.86698483e-04,
                    8.98489491e-05,
                    6.95804694e-05,
                    1.00709658e-04,
                    6.71267555e-05,
                    6.21248660e-05,
                    9.77974090e-05,
                    5.55782024e-05,
                    7.86778001e-05,
                ],
                [
                    2.15052791e-04,
                    2.57932095e-04,
                    1.58958119e-04,
                    1.53111000e-04,
                    1.28918711e-04,
                    1.28399776e-04,
                    1.52606397e-04,
                    8.74914106e-05,
                    1.45795307e-04,
                    8.58786148e-05,
                    1.11349186e-04,
                    8.98489491e-05,
                    2.95359997e-04,
                    1.03311880e-04,
                    9.29874364e-05,
                    9.31063812e-05,
                    1.25233072e-04,
                    1.38130001e-04,
                    8.47903189e-05,
                    1.02860310e-04,
                ],
                [
                    9.97625137e-05,
                    1.02358608e-04,
                    9.03141055e-05,
                    8.80192972e-05,
                    8.25067070e-05,
                    7.87550632e-05,
                    1.00653474e-04,
                    7.90839398e-05,
                    8.85816243e-05,
                    1.00394611e-04,
                    8.26738986e-05,
                    6.95804694e-05,
                    1.03311880e-04,
                    1.38679551e-04,
                    7.33972243e-05,
                    9.85266422e-05,
                    5.69635333e-05,
                    9.18906352e-05,
                    7.76157856e-05,
                    7.37991300e-05,
                ],
                [
                    8.61115392e-05,
                    1.10805178e-04,
                    1.04095672e-04,
                    7.46460062e-05,
                    8.71845732e-05,
                    9.10274252e-05,
                    8.12084880e-05,
                    8.87031919e-05,
                    1.00943437e-04,
                    6.65754039e-05,
                    1.18642453e-04,
                    1.00709658e-04,
                    9.29874364e-05,
                    7.33972243e-05,
                    2.01771356e-04,
                    6.80103047e-05,
                    8.57837055e-05,
                    1.03111611e-04,
                    5.68981153e-05,
                    7.80565507e-05,
                ],
                [
                    8.56846596e-05,
                    8.74489219e-05,
                    8.01562125e-05,
                    8.14494768e-05,
                    6.50007767e-05,
                    7.15203459e-05,
                    8.22072630e-05,
                    7.56739267e-05,
                    7.67240024e-05,
                    8.42955525e-05,
                    7.54084960e-05,
                    6.71267555e-05,
                    9.31063812e-05,
                    9.85266422e-05,
                    6.80103047e-05,
                    1.36919070e-04,
                    6.54501975e-05,
                    8.09845020e-05,
                    7.37370174e-05,
                    5.87843464e-05,
                ],
                [
                    1.37608444e-04,
                    2.68568733e-04,
                    2.67054571e-04,
                    2.06813151e-04,
                    3.23877738e-04,
                    2.41363768e-04,
                    1.01272428e-04,
                    6.55688622e-05,
                    2.16800171e-04,
                    5.75321383e-05,
                    7.99199098e-05,
                    6.21248660e-05,
                    1.25233072e-04,
                    5.69635333e-05,
                    8.57837055e-05,
                    6.54501975e-05,
                    1.46893079e-03,
                    9.85348107e-05,
                    6.59574929e-05,
                    3.03643568e-04,
                ],
                [
                    1.33035904e-04,
                    1.62238509e-04,
                    1.52271435e-04,
                    1.25016166e-04,
                    1.36417816e-04,
                    1.24911263e-04,
                    1.24468155e-04,
                    9.38343175e-05,
                    1.38853956e-04,
                    8.11275129e-05,
                    1.17689173e-04,
                    9.77974090e-05,
                    1.38130001e-04,
                    9.18906352e-05,
                    1.03111611e-04,
                    8.09845020e-05,
                    9.85348107e-05,
                    2.64604165e-04,
                    7.18645008e-05,
                    1.09645039e-04,
                ],
                [
                    8.21656196e-05,
                    8.62748168e-05,
                    6.72958419e-05,
                    9.30891364e-05,
                    5.45391833e-05,
                    5.27462412e-05,
                    8.57347703e-05,
                    6.01344712e-05,
                    6.40983892e-05,
                    5.74148452e-05,
                    6.39652089e-05,
                    5.55782024e-05,
                    8.47903189e-05,
                    7.76157856e-05,
                    5.68981153e-05,
                    7.37370174e-05,
                    6.59574929e-05,
                    7.18645008e-05,
                    1.79889721e-04,
                    5.13016897e-05,
                ],
                [
                    1.06281426e-04,
                    1.38462994e-04,
                    1.98387376e-04,
                    1.27462037e-04,
                    2.79531990e-04,
                    1.95714717e-04,
                    9.68811805e-05,
                    7.14216054e-05,
                    1.77534524e-04,
                    8.46645084e-05,
                    8.04271042e-05,
                    7.86778001e-05,
                    1.02860310e-04,
                    7.37991300e-05,
                    7.80565507e-05,
                    5.87843464e-05,
                    3.03643568e-04,
                    1.09645039e-04,
                    5.13016897e-05,
                    3.08721621e-04,
                ],
            ]
        ),
    )
    views = ["G1 - G2 == 0.03 ", "G3 - G4== 0.04", "G5 == 0.06 "]
    groups = {
        "AAPL": ["G1"],
        "BBY": ["G2"],
        "CVX": ["G3"],
        "KO": ["G4"],
        "MSFT": ["G5"],
    }
    model2 = BlackLitterman(views=views, groups=groups, tau=1 / n_observations)
    model2.fit(X)
    np.testing.assert_almost_equal(
        model2.return_distribution_.mu, model.return_distribution_.mu
    )
    np.testing.assert_almost_equal(
        model2.return_distribution_.covariance, model.return_distribution_.covariance
    )


def test_black_litterman_factor_model(X, y):
    views = ["AAPL - BBY == 0.03 ", "MSFT == 0.06 "]
    factor_views = ["MTUM - QUAL == 0.03 ", "VLUE == 0.06"]

    model = BlackLitterman(
        views=views,
        prior_estimator=FactorModel(
            factor_prior_estimator=BlackLitterman(views=factor_views),
        ),
    )

    model.fit(X, y)
    assert model.return_distribution_.covariance.shape == (20, 20)


def test_metadata_routing(X, implied_vol):
    views = ["AAPL - BBY == 0.03 ", "MSFT == 0.06 "]

    with config_context(enable_metadata_routing=True):
        model = BlackLitterman(
            views=views,
            prior_estimator=EmpiricalPrior(
                covariance_estimator=ImpliedCovariance().set_fit_request(
                    implied_vol=True
                )
            ),
        )

        with pytest.raises(ValueError):
            model.fit(X)

        model.fit(X, implied_vol=implied_vol)

    # noinspection PyUnresolvedReferences
    assert model.prior_estimator_.covariance_estimator_.r2_scores_.shape == (20,)
